import cv2
import numpy as np
from typing import Optional
import time
import threading
from queue import Queue
from rknnpool import rknnPoolExecutor
import logging
logging.basicConfig(level=logging.WARNING)
from image_provider import ImageProvider
from logger import setup_logger
from seg_inference2 import seg_func

# 初始化日志
logger = setup_logger()
segmodelPath = "./models/merge_seg_int8_rk3588.rknn"
posemodelPath = "./models/merge_yolov11_pose.rknn"

class ImageData:
    """图像数据结构"""
    def __init__(self):
        self.timestamp: float = 0.0
        self.readable_timestamp: str = ""
        self.color_img: np.ndarray = None
        self.depth_img: np.ndarray = None
        self.color_path: Optional[str] = None

class AsyncProcessor:
    def __init__(self, seg_pool, target_fps=30):
        self.seg_pool = seg_pool
        self.target_fps = target_fps
        self.running = False
        self.image_provider = ImageProvider(from_folder=True)
        
        # 线程相关变量
        # 帧率统计
        self.submit_count = 0  # 提交帧计数
        self.process_count = 0  # 处理完成计数
        self.last_log_time = time.time()
        
        # 性能计时
        self.total_inference_time = 0  # 总推理时间(秒)
        self.total_frames_processed = 0  # 总处理帧数
        
    def start(self):
        """启动异步处理器"""
        self.running = True
        # 启动生产者线程（以固定速率提交帧）
        self.producer_thread = threading.Thread(target=self._producer)
        self.producer_thread.daemon = True  # 设置为守护线程
        self.producer_thread.start()
        
        # 启动消费者线程（处理结果）
        self.consumer_thread = threading.Thread(target=self._consumer)
        self.consumer_thread.daemon = True  # 设置为守护线程
        self.consumer_thread.start()
    
    def stop(self):
        """停止异步处理器"""
        self.running = False
        self.producer_thread.join(timeout=1)
        self.consumer_thread.join(timeout=1)
    
    def _producer(self):
        """生产者线程：尝试以固定帧率提交帧"""
        frame_interval = 1.0 / self.target_fps
        next_frame_time = time.time()
        
        while self.running:
            start_time = time.time()
            
            image = self.image_provider.get_next_image()
            if image is None:
                break
                
            # 提交帧到处理池
            self.seg_pool.put(image.color_img)
            self.submit_count += 1
            
            # 计算需要睡眠的时间以维持目标FPS
            processing_time = time.time() - start_time
            sleep_time = max(0, frame_interval - processing_time)
            time.sleep(sleep_time/2)
            
            # 更新下一帧时间
            next_frame_time += frame_interval
            
            # 每秒记录一次帧率
            self._log_fps()
    
    def _consumer(self):
        """消费者线程：处理结果并计算推理速度"""
        while self.running:
            start_time = time.time()
            
            # 获取处理结果
            seg_result, flag = self.seg_pool.get()
                
            # 计算推理耗时
            inference_time = time.time() - start_time
            self.total_inference_time += inference_time
            if flag:
                self.total_frames_processed += 1
                self.process_count += 1
            
            # 每秒记录一次帧率
            self._log_fps()
    
    def _log_fps(self):
        """记录帧率信息"""
        current_time = time.time()
        if current_time - self.last_log_time >= 1.0:
            elapsed = current_time - self.last_log_time
            
            # 计算实际帧率
            submit_fps = self.submit_count / elapsed
            process_fps = self.process_count / elapsed
            
            # 计算平均推理时间
            avg_inference = 0
            if self.total_frames_processed > 0:
                avg_inference = self.total_inference_time / self.total_frames_processed
            
            logger.info(
                f"帧率统计: 提交={submit_fps:.1f}FPS, "
                f"处理={process_fps:.1f}FPS, "
                f"延迟={avg_inference*1000:.1f}ms/帧 "
                f"(目标={self.target_fps}FPS)"
            )
            
            # 重置计数器
            self.submit_count = 0
            self.process_count = 0
            self.last_log_time = current_time

# 初始化线程池
TPEs = 5
logger.info(f"初始化RKNN线程池，线程数: {TPEs}")
seg_pool = rknnPoolExecutor(
    rknnModel=segmodelPath,
    TPEs=TPEs,
    func=seg_func,
    npus=[1,2]
)

# 创建并启动异步处理器
logger.info("启动异步处理器，目标帧率: 30 FPS")
async_processor = AsyncProcessor(seg_pool, target_fps=30)
async_processor.start()

# 主线程可以在这里做其他工作或只是等待
try:
    logger.info("主线程运行中，按Ctrl+C退出...")
    while True:
        time.sleep(1)
except KeyboardInterrupt:
    logger.info("接收到中断信号，正在停止...")
    async_processor.stop()
    logger.info("程序已安全停止")
finally:
    cv2.destroyAllWindows()